#include <stdlib.h>
#include <cuda_runtime.h>
#include <cublas_v2.h>
#include <time.h>
#include <iostream>
#include <functional>
#include <algorithm>
#include "utils/macro.h"
#include "utils/timer_guard.h"
#include "kernels/sgemm_shared_mem_block.cuh"
#include "kernels/sgemm_BM_BN_BK_TM.cuh"
#include "kernels/sgemm_BM_BN_BK_TM_TN.cuh"
#include "kernels/sgemm_BM_BN_BK_TM_TN_vertorize.cuh"

size_t M = 4096 + 7;
size_t N = 4096 + 9;
size_t K = 4096 + 11;
#define NUM_THREADS 1024

__host__ void init_matrix(float* m, size_t n) {
    for (size_t i = 0; i < n; i ++) {
        m[i] = float(rand() % 110) / 55 - 1.0f;
    }
}

#define TEST_SGEMM_KERNEL(module, gdim, bdim)                                                       \
    {                                                                                               \
        std::cout << #module << ": " << std::endl;                                                  \
        float* tmpc;                                                                                \
        cudaMalloc((void**)&tmpc, M * N * sizeof(float));                                           \
        {                                                                                           \
            TIME_GUARD(TEST_SGEMM_KERNEL);                                                          \
            module<<<gdim, bdim>>>(M, N, K, alpha, a, b, beta, tmpc);                               \
            CUDA_CHECK_STATUS();                                                                    \
            cudaDeviceSynchronize();                                                                \
        }                                                                                           \
        CUDA_CHECK_STATUS();                                                                        \
        cudaMemcpy(thc, tmpc, M * N * sizeof(float), cudaMemcpyDeviceToHost);                       \
        VALUE_CHECK(rhc, thc, (M * N), module);                                                     \
        cudaFree(tmpc);                                                                             \
        std::cout << std::endl;                                                                     \
    }

#define test_shared_mem_block(BLOCK_SIZE)                                               \
    {                                                                                   \
        dim3 grid(CEIL_DIV(M, BLOCK_SIZE), CEIL_DIV(N, BLOCK_SIZE));                    \
        dim3 block(BLOCK_SIZE * BLOCK_SIZE);                                            \
        auto sgemm_shared_mem_block_##BLOCK_SIZE = sgemm_shared_mem_block<BLOCK_SIZE>;  \
        TEST_SGEMM_KERNEL(sgemm_shared_mem_block_##BLOCK_SIZE, grid, block);            \
    }

#define test_BM_BN_BK_TM(BM, BK)                                                        \
    {                                                                                   \
        const int BN = BM;                                                              \
        const int TM = BM * BN / (BM * BK);                                             \
        dim3 grid(CEIL_DIV(M, BM), CEIL_DIV(N, BN));                                    \
        dim3 block(BM * BK);                                                            \
        auto sgemm_BM_BN_BK_TM_##BM##_##BK = sgemm_BM_BN_BK_TM<BM, BN, BK, TM>;         \
        TEST_SGEMM_KERNEL(sgemm_BM_BN_BK_TM_##BM##_##BK, grid, block);                  \
    }

#define test_BM_BN_BK_TM_TN(BM, BK, TM, TN)                                                                 \
    {                                                                                                       \
        const int BN = BM;                                                                                  \
        dim3 grid(CEIL_DIV(M, BM), CEIL_DIV(N, BN));                                                        \
        dim3 block(BM * BN / (TM * TN));                                                                    \
        auto sgemm_BM_BN_BK_TM_TN_##BM##_##BK##_##TM##_##TN = sgemm_BM_BN_BK_TM_TN<BM, BN, BK, TM, TN>;     \
        TEST_SGEMM_KERNEL(sgemm_BM_BN_BK_TM_TN_##BM##_##BK##_##TM##_##TN, grid, block);                     \
    }

int main() {
    srand(time(NULL));
    float alpha = 1.0f;
    float beta = 0.0f;
    cublasHandle_t handle;
    cublasCreate(&handle);

    float *a, *b, *c;
    float *ha = (float*)malloc(M * K * sizeof(float));
    float *hb = (float*)malloc(K * N * sizeof(float));
    float *rhc = (float*)malloc(M * N * sizeof(float));
    float *thc = (float*)malloc(M * N * sizeof(float));
    
    cudaMalloc((void**)&a, M * K * sizeof(float));
    cudaMalloc((void**)&b, K * N * sizeof(float));
    cudaMalloc((void**)&c, M * N * sizeof(float));

    {
        TIME_GUARD(INIT);
        init_matrix(ha, M * K);
        init_matrix(hb, K * N);
        cudaMemcpy(a, ha, M * K * sizeof(float), cudaMemcpyHostToDevice);
        cudaMemcpy(b, hb, K * N * sizeof(float), cudaMemcpyHostToDevice);
        CUDA_CHECK_STATUS();
    }

    cudaDeviceSynchronize();
    cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, N, M, K, &alpha, b, N, a, K, &beta, c, N);
    cudaDeviceSynchronize();
    {
        TIME_GUARD(CUBLAS);
        cublasSgemm(handle, CUBLAS_OP_N, CUBLAS_OP_N, N, M, K, &alpha, b, N, a, K, &beta, c, N);
        cudaDeviceSynchronize();
        CUDA_CHECK_STATUS();
    }
    cudaMemcpy(rhc, c, M * N * sizeof(float), cudaMemcpyDeviceToHost);

    test_shared_mem_block(8);
    test_shared_mem_block(16);
    test_shared_mem_block(32);

    test_BM_BN_BK_TM(16, 8);
    test_BM_BN_BK_TM(32, 8);
    test_BM_BN_BK_TM(64, 8);
    test_BM_BN_BK_TM(64, 16);

    test_BM_BN_BK_TM_TN(64, 8, 4, 4);
    test_BM_BN_BK_TM_TN(64, 8, 8, 4);
    test_BM_BN_BK_TM_TN(64, 8, 8, 8);
    test_BM_BN_BK_TM_TN(128, 8, 8, 8);

    {
        const int BM = 64, BN = 64, BK = 8, TM = 8, TN = 4;
        dim3 grid(CEIL_DIV(M, BM), CEIL_DIV(N, BN));
        dim3 block(BM * BN / (TM * TN));

        auto sgemm_128_128_8_8_8_vertorize = sgemm_BM_BN_BK_TM_TN_vertorize<BM, BN, BK, TM, TN>;
        TEST_SGEMM_KERNEL(sgemm_128_128_8_8_8_vertorize, grid, block);
    }

    cudaFree(a);
    cudaFree(b);
    cudaFree(c);
    free(ha);
    free(hb);
    free(rhc);
    free(thc);

    return 0;
}